package com.zws.cucumber.adaptation.cass;

import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.cql.*;
import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
import com.datastax.oss.driver.api.core.type.codec.TypeCodec;
import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
import com.datastax.oss.driver.internal.core.type.codec.StringCodec;
import com.datastax.oss.driver.internal.core.util.Strings;
import com.zws.cucumber.util.colorfullog.ZwsColorfulLogger;
import edu.umd.cs.findbugs.annotations.NonNull;
import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

import static com.zws.cucumber.adaptation.cass.ZwsCqlDmlType.Insert;
import static java.util.Objects.requireNonNull;

public class ZwsCqlSessionProxyHelper {

    private static final Logger logger = ZwsColorfulLogger.of(ZwsCqlSessionProxyHelper.class);

    /*
     * tableName -> metadata
     * tableName -> delete PreparedStmts
     */
    private static CodecRegistry codecRegistry;
    private static final ConcurrentMap<String, TableMetadata> tableMetadataMap = new ConcurrentHashMap<>(256);
    private static final ConcurrentMap<String, List<ColumnMetadata>> partitionKeyMetadatasMap = new ConcurrentHashMap<>(256);
    private static final ConcurrentMap<String, PreparedStatement> deletePreparedStmts = new ConcurrentHashMap<>(256);

    private static final ConcurrentMap<String, PreparedStatement> preparedStmtMap = new ConcurrentHashMap<>();
    private static final ConcurrentMap<PreparedStatement, ZwsCqlInfo> preparedStmt2CqlInfoMap = new ConcurrentHashMap<>();

    // tableName -> partitionKeyAndValueList
    private static final ConcurrentMap<String, List<Map<String, Object>>> dirtyDataMap = new ConcurrentHashMap<>();

    private static ExecutorService executorService = new ThreadPoolExecutor(32, 32, 1, TimeUnit.MINUTES, new ArrayBlockingQueue<>(128));

    private static volatile Future<?> asyncCleanDataFuture = null;

    @Deprecated
    public static void init0(CqlSession session) {
        codecRegistry = session.getContext().getCodecRegistry();

        Collection<TableMetadata> tableMetadatas = session.getMetadata()
                .getKeyspace(session.getKeyspace().get())
                .get()
                .getTables()
                .values();
        logger.debug("tableMetadatas.size: {}", tableMetadatas.size());

        CountDownLatch latch = new CountDownLatch(tableMetadatas.size());
        for (TableMetadata metadata : tableMetadatas) {
            String tableName = metadata.getName().asInternal();
            tableMetadataMap.put(tableName, metadata);
            List<ColumnMetadata> partitionKeyMetadatas = metadata.getPartitionKey();
            partitionKeyMetadatasMap.put(tableName, partitionKeyMetadatas);
            // delete
            executorService.submit(() -> {
                logger.debug("session keyspace: {}", session.getKeyspace().get());
                PreparedStatement preparedStmt = session.prepare(
                    generateDeletePreparedCql(tableName, partitionKeyMetadatas));
                deletePreparedStmts.put(tableName, preparedStmt);
                latch.countDown();
            });
        }

        waitLatch(latch);

        System.out.println(tableMetadataMap);
        System.out.println(partitionKeyMetadatasMap);
        System.out.println(deletePreparedStmts);
    }

    public static void init1(CqlSession session) {
        long s = System.currentTimeMillis();

        codecRegistry = session.getContext().getCodecRegistry();

        Collection<TableMetadata> tableMetadatas = session.getMetadata()
                .getKeyspace(session.getKeyspace().get())
                .get()
                .getTables()
                .values();

        for (TableMetadata metadata : tableMetadatas) {
            String tableName = metadata.getName().asInternal();
            tableMetadataMap.put(tableName, metadata);
        }

        /*
         * not generate prepare statement, it seems no uses.
         */
//        CountDownLatch latch = new CountDownLatch(tableMetadatas.size());
//        for (TableMetadata metadata : tableMetadatas) {
//            String tableName = metadata.getName().asInternal();
//            // generate insert prepared statement
//            executorService.submit(() -> {
//                try {
//                    long ss = System.currentTimeMillis();
//                    prepareInsertStatement(session, tableName);
//                    logger.debug("generate insertion prepared statement for {} finished, elapsed: {}ms", tableName, System.currentTimeMillis() - ss);
//                } catch (Exception e) {
//                    logger.error("prepareInsertStatement failed", e);
//                } finally {
//                    latch.countDown();
//                }
//            });
//        }
//        waitLatch(latch);
        logger.debug("ZwsCqlSessionProxyHelper.init() finished, elapsed: {}ms", System.currentTimeMillis() - s);
    }

    private static String generateDeletePreparedCql(String tableName, List<ColumnMetadata> partitionKeyMetadataList) {
        String singleColumnCqlTemplate = "delete from %s where %s in :l";
        String multipleColumnCqlTemplate = "delete from %s where %s";
        String cql;
        if (partitionKeyMetadataList.size() == 1) {
            cql = String.format(singleColumnCqlTemplate, tableName, partitionKeyMetadataList.get(0).getName().asInternal());
        } else {
            List<String> partitionKeyList = partitionKeyMetadataList.stream()
                    .map(m -> m.getName().asInternal())
                    .collect(Collectors.toList());
            List<String> conditions = new ArrayList<>();
            for (int i = 0; i < partitionKeyList.size(); i++) {
                conditions.add(partitionKeyList.get(i) + "=:" + partitionKeyList.get(i));
            }
            cql = String.format(multipleColumnCqlTemplate, tableName, String.join(" and ", conditions));
        }
        logger.debug("generateDeletePreparedCql: {}", cql);
        return cql;
    }

    public static void addDirtyData(String tableName, Map<String, Object> partitionKeysAndValues) {
        requireNonNull(partitionKeysAndValues);
        List<Map<String, Object>> list = dirtyDataMap.computeIfAbsent(tableName, k -> new ArrayList<>());
        list.add(partitionKeysAndValues);
    }

    public static <RequestT> void intercept(@NonNull RequestT request) {
//        logger.debug("intercept, request: {}", request);
        ZwsCqlInfo info = null;
        if (request instanceof String) {
            info = (ZwsCqlExtendedInfo) ZwsCqlExtractor.extract((String) request, codecRegistry);
            if (info != null) {
                addDirtyData(info.getTableName(), ((ZwsCqlExtendedInfo) info).getPartitionKeyAndValueMap());
            }

        } else if (request instanceof SimpleStatement) {
            info = (ZwsCqlExtendedInfo) ZwsCqlExtractor.extract(((SimpleStatement) request).getQuery(), codecRegistry);
            if (info != null) {
                addDirtyData(info.getTableName(), ((ZwsCqlExtendedInfo) info).getPartitionKeyAndValueMap());
            }

        } else if (request instanceof BoundStatement) {
            BoundStatement boundStmt = (BoundStatement) request;
            PreparedStatement preparedStmt = boundStmt.getPreparedStatement();
            info = preparedStmt2CqlInfoMap.get(preparedStmt);
            if (info == null) {
                logger.debug("extract ZwsCqlInfo due to prepared statement not in preparedStmt2CqlInfoMap");
                preparedStmt2CqlInfoMap.putIfAbsent(
                        preparedStmt, info = Optional.ofNullable(ZwsCqlExtractor.extract(preparedStmt.getQuery())).orElse(ZwsCqlInfo.NULL));
            }

            if (info != ZwsCqlInfo.NULL) {
                addDirtyData(info.getTableName(), buildPartitionKeysAndValues(boundStmt));
            }

        } else if (request instanceof BatchStatement) {
            BatchStatement batchStmt = (BatchStatement) request;
            for (BatchableStatement<?> stmt : batchStmt) {
                intercept(stmt);
            }
        } else {
            logger.debug("skip request: {}", request);
        }
        logger.trace("after intercept, info: {}", info);
        logger.trace("after intercept, dirty data map: {}", dirtyDataMap);
    }

    private static Map<String, Object> buildPartitionKeysAndValues(BoundStatement boundStmt) {
        PreparedStatement preparedStmt = boundStmt.getPreparedStatement();
        List<Integer> partitionKeyIndices = preparedStmt.getPartitionKeyIndices();
        if (CollectionUtils.isEmpty(partitionKeyIndices)) {
            throw new UnsupportedOperationException();
        }
        Map<String, Object> partitionKeysAndValues = new HashMap<>();
        for (Integer idx : partitionKeyIndices) {
            ColumnDefinition def = preparedStmt.getVariableDefinitions().get(idx);
            partitionKeysAndValues.put(
                    def.getName().asInternal(),
                    boundStmt.get(idx, codecRegistry.codecFor(def.getType())));
        }
        return partitionKeysAndValues;
    }

    public static TableMetadata getTableMetadata(String tableName) {
        if (tableName.contains(".")) {
            tableName = tableName.substring(tableName.indexOf('.') + 1);
        }
        return tableMetadataMap.get(tableName);
    }

    public static void cleanData(CqlSession session, boolean async) {
        final long startTimestamp = System.currentTimeMillis();
        CountDownLatch latch = new CountDownLatch(dirtyDataMap.size());
        logger.debug("clean data in table: {}", dirtyDataMap.keySet());
        logger.trace("dirtyDataMap: {}", dirtyDataMap);
        for (Map.Entry<String, List<Map<String, Object>>> entry : dirtyDataMap.entrySet()) {
            String tableName = entry.getKey();
            List<Map<String, Object>> partitionKeysAndValuesList = entry.getValue();
            PreparedStatement preparedStmt = deletePreparedStmts.get(tableName);
            logger.trace("prepared statement: {}, cql: {}", preparedStmt, preparedStmt.getQuery());
            requireNonNull(preparedStmt, "no prepared statement for table: " + tableName);
            TableMetadata tableMetadata = tableMetadataMap.get(tableName);
            requireNonNull(tableMetadata, "no table metadata for table: " + tableName);
            Statement[] stmtHolder = new Statement[1];
            logger.trace("bound statement, partitionKeysAndValuesList: {}", partitionKeysAndValuesList);

            if (tableMetadata.getPartitionKey().size() == 1) {
                String partitionKey = tableMetadata.getPartitionKey().get(0).getName().asInternal();
                logger.trace("bound statement, partitionKey: {}", partitionKey);
                List values = partitionKeysAndValuesList.stream()
                        .map(m -> m.get(partitionKey))
                        .collect(Collectors.toList());
                logger.trace("bound statement, value list: {}", values);
                stmtHolder[0] = preparedStmt.bind().setList(0, values, values.get(0).getClass());

            } else {
                BatchStatement batchStmt = BatchStatement.newInstance(BatchType.LOGGED);
                for (Map<String, Object> partitionKeysAndValues : partitionKeysAndValuesList) {
                    BoundStatement boundStmt = preparedStmt.bind();
                    for (ColumnMetadata partitionKeyMetadata : tableMetadata.getPartitionKey()) {
                        String key = partitionKeyMetadata.getName().asInternal();
                        Object val = partitionKeysAndValues.get(key);
                        boundStmt = boundStmt.set(key, val, (Class) val.getClass());
                    }
                    batchStmt = batchStmt.add(boundStmt);
                }
                stmtHolder[0] = batchStmt;
            }
            executorService.submit(() -> {
                // execute it
                session.execute(stmtHolder[0]);
                latch.countDown();
            });

        }
        if (async) {
            asyncCleanDataFuture = executorService.submit(() -> {
                waitLatch(latch);
                dirtyDataMap.clear();
                logger.debug("async clean data done, elapsed: {}ms", System.currentTimeMillis() - startTimestamp);
            });
            logger.debug("submit async clean data task, elapsed: {}ms", System.currentTimeMillis() - startTimestamp);
        } else {
            waitLatch(latch);
            dirtyDataMap.clear();
            logger.debug("sync clean data done, elapsed: {}ms", System.currentTimeMillis() - startTimestamp);
        }

    }

    private static void prepareInsertStatement(CqlSession session, String tableName) {
        TableMetadata metadata = tableMetadataMap.get(tableName);
        Objects.requireNonNull(metadata, "not existed metadata for table: " + tableName);
        // one for primary key
        // one for all column
        List<Collection<ColumnMetadata>> list = new ArrayList<>();
        list.add(metadata.getPrimaryKey());
        list.add(metadata.getColumns().values());
        for (Collection<ColumnMetadata> columnMetadataList : list) {
            List<String> columnNameList = columnMetadataList.stream()
                    .map(md -> md.getName().asInternal())
                    .sorted()
                    .collect(Collectors.toList());

            String key = getPreparedStatementKey(tableName, Insert, columnNameList);
            String cql = generateInsertionPreparedCql(tableName, columnNameList);
            logger.debug("generate insertion prepared statement, key: {}, cql: {}", key, cql);
            PreparedStatement preparedStmt = session.prepare(cql);
            preparedStmtMap.put(key, preparedStmt);
        }
    }

    private static String getPreparedStatementKey(String tableName, ZwsCqlDmlType dmlType, Collection<String> columnNames) {
        StringJoiner joiner = new StringJoiner("-");
        joiner.add(tableName)
                .add(dmlType.name());
        for (String columnName : columnNames) {
            joiner.add(columnName);
        }
        return joiner.toString();
    }

    private static String generateInsertionPreparedCql(String tableName, Collection<String> columnNames) {
        return new StringJoiner(" ")
                .add("insert")
                .add("into")
                .add(tableName)
                .add("(")
                .add(String.join(", ", columnNames))
                .add(")")
                .add("values")
                .add("(")
                .add(columnNames.stream().map(c -> ":" + c).collect(Collectors.joining(", ")))
                .add(")")
                .toString();
    }

    private static Collection<String> sortColumnNames(Collection<String> columnNames) {
        return columnNames.stream()
                .sorted()
                .collect(Collectors.toList());
    }

    public static PreparedStatement getInsertionPreparedStatement(CqlSession session, String tableName, Collection<String> columnNames) {
        Collection<String> sortedColumnNames = sortColumnNames(columnNames);
        String key = getPreparedStatementKey(tableName, Insert, sortedColumnNames);
        return preparedStmtMap.computeIfAbsent(key, k -> {
            long s = System.currentTimeMillis();
            String cql = generateInsertionPreparedCql(tableName, sortedColumnNames);
            logger.debug("generate insertion prepared statement, elapsed: {}ms, key: {}, cql: {}",
                    System.currentTimeMillis() - s, key, cql);
            return session.prepare(cql);
        });
    }

    public static BoundStatement generateBoundStatement(TableMetadata metadata, PreparedStatement preparedStmt, Map<String, String> row) {
        BoundStatementBuilder stmtBuilder = preparedStmt.boundStatementBuilder();

        row.forEach((k, v) -> {
            TypeCodec codec = codecRegistry.codecFor(metadata.getColumn(k).get().getType());
            if (codec instanceof StringCodec && v != null) {
                v = Strings.quote(v);
            }
            stmtBuilder.set(k, codec.parse(v), codec);
        });
        return stmtBuilder.build();
    }

    private static void waitLatch(CountDownLatch latch) {
        try {
            latch.await();
        } catch (InterruptedException e) {
            logger.error("", e);
        }
    }

    public static void waitAsyncCleanData() {
        if (asyncCleanDataFuture != null) {
            try {
                asyncCleanDataFuture.get();
            } catch (Exception e) {
                throw new RuntimeException(e);
            } finally {
                asyncCleanDataFuture = null;
            }
        }
    }

}
